import sys
import os
import cv2
import matplotlib.pyplot as plt

from ros import rosbag
import cv_bridge
import rospy

from sensor_msgs.msg import Imu
from geometry_msgs.msg import Vector3
from geometry_msgs.msg import Point
from geometry_msgs.msg import PointStamped

image_path = '/home/yujr/rosbag-tool/saveBag/downshort/image'
image_list = os.listdir(image_path)
image_list.sort()

bag = rosbag.Bag('downshort.bag','w')
bridge = cv_bridge.CvBridge()


for image_name in image_list:
    full_path = image_path + '/' + image_name
    img_original = cv2.imread(full_path)
    # print(img_original.shape)
    # img_crop=img_original[:800,:1900,:]
    # print(img_crop.shape)
    img=cv2.resize(img_original,(752,480),interpolation=cv2.INTER_CUBIC)
    # print(img.shape)
    cv2.imshow('image', img)
    cv2.waitKey(30)
    timestamp = os.path.splitext(image_name)[0]
    
    img_msg = bridge.cv2_to_imgmsg(img)
    img_msg.encoding = 'rgb8'
    
    stamp = rospy.rostime.Time.from_sec(float(timestamp))
    bag.write('/cam0/image_raw', img_msg, stamp)
    bag.write('/cam1/image_raw', img_msg, stamp)


imu_file = '/home/yujr/rosbag-tool/saveBag/downshort/imu.txt'
imuFile = open(imu_file,'r')
all = imuFile.readlines()
imu_timestamp = []
imu_data = []
for f in all:
    line = f.rstrip('\n').split(' ')
    imu_timestamp.append(line[0])
    imu_data.append(line[1:])

for i in range(len(imu_data)):
    imu = Imu()
    
    angular_v=Vector3()
    linear_a=Vector3()
    
    linear_a.x=float(imu_data[i][0])
    linear_a.y=float(imu_data[i][1])
    linear_a.z=float(imu_data[i][2])
    angular_v.x=float(imu_data[i][3])
    angular_v.y=float(imu_data[i][4])
    angular_v.z=float(imu_data[i][5])
    
    # print(linear_a)
    # print(angular_v)
    # print(imu_timestamp[i])
    # print(float(imu_timestamp[i]))
    
    imuStamp = rospy.rostime.Time.from_sec(float(imu_timestamp[i])/1e3)
    
    print(imuStamp)
    
    imu.header.stamp=imuStamp
    imu.angular_velocity = angular_v
    imu.linear_acceleration = linear_a
    bag.write("/imu0",imu,imuStamp)


ptStampeFile = '/home/yujr/rosbag-tool/saveBag/downshort/associate.txt'
fpt = open(ptStampeFile,"r")
data = fpt.readlines()
fpt.close()
vt=[]
vx=[]
vy=[]
vz=[]
for i in range(len(data)):
    tmp = data[i].split(' ')
    vt.append(tmp[0])
    vx.append(tmp[1])
    vy.append(tmp[2])
    vz.append(tmp[3])

for i in range(len(vt)):
    pt=PointStamped()

    ptStamp = rospy.rostime.Time.from_sec(float(vt[i])/1e3)
    print(ptStamp)
    pt.header.stamp=ptStamp

    ptxyz=Point()
    ptxyz.x=vx[i]
    ptxyz.y=vy[i]
    ptxyz.z=vz[i]
    pt.point=ptxyz
    bag.write("/leica/position",pt,ptStamp)

bag.close()
